/*
 * Copyright (C) 2008 Nokia Corporation
 * Copyright (c) International Business Machines Corp., 2006
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See
 * the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/*
 * Generate UBI images.
 *
 * Authors: Artem Bityutskiy
 *          Oliver Lohmann
 */

#include <stdio.h>
#include <sys/stat.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <Windows.h>
#include <io.h>

#include "libubi.h"
#include "ubi-media.h"
#include "libubigen.h"
#include "ubifs-media.h"
#include "ubifs.h"
#include "unubinize.h"
#include "utils.h"

static int read_volume_info(const struct ubigen_info *ui,
		struct volume_info_node *vi_node, struct ubigen_vol_info *vi,
		struct stat *st)
{
	const char *vol_name = vi_node->name;
	const char *img = vi_node->image_path;

	/* Fetch volume type */
	vi->type = vi_node->vol_type;
	if (vi->type != UBI_VID_STATIC && vi->type != UBI_VID_DYNAMIC)
		return errmsg("invalid volume type %d in volume \"%s\"\n",  vi->type, vol_name);

	if (verbose)
		infomsg("volume type: %s\n", vi->type == UBI_VID_DYNAMIC ? "dynamic" : "static");

	/* Fetch the name of the volume image file */
	if (img) {
		if (stat(img, st))
			return errmsg("cannot stat \"%s\" referred from volume \"%s\"\n",
					  img, vol_name);
		if (st->st_size == 0)
			return errmsg("empty file \"%s\" referred from volume \"%s\"\n",
				       img, vol_name);
	} else if (vi->type == UBI_VID_STATIC)
		return errmsg("image is not specified for static volume \"%s\"\n",
		               vol_name);

	/* Fetch volume id */
	vi->id = vi_node->vol_id;
	if (vi->id == -1)
		return errmsg("\"vol_id\" is not valid in volume \"%s\"\n", vol_name);
	if (vi->id < 0)
		return errmsg("negative volume ID %d in section \"%s\"\n",
			      vi->id, vol_name);
	if (vi->id >= ui->max_volumes)
		return errmsg("too high volume ID %d in section \"%s\", max. is %d\n",
			      vi->id, vol_name, ui->max_volumes);

	if (verbose)
		infomsg("volume ID: %d\n", vi->id);

	/* Fetch volume size */
	if (vi_node->reserved_pebs > 0) {
		vi->bytes = vi_node->reserved_pebs * ui->leb_size;   /* only lebs contain file data */
		if (vi->bytes <= 0)
			return errmsg("bad vol_size \"%d\" in volume \"%s\"\n", (int)vi->bytes, vol_name);

		/* Make sure the image size is not larger than volume size */
		if (img && st->st_size > vi->bytes)
			return errmsg("error in volume \"%s\": size of the image file "
						"\"%s\" is %lld, which is larger than volume size %lld\n",
						vol_name, img, (long long)st->st_size, vi->bytes);
		if (verbose)
			infomsg("volume size: %lld bytes\n", vi->bytes);
	} else {
		if (!img)
			return errmsg("neither image file (\"image=\") nor volume size "
				      "(\"vol_size=\") specified in volume \"%s\"\n", vol_name);

		vi->bytes = st->st_size;

		if (vi->bytes == 0)
			return errmsg("file \"%s\" referred from volume \"%s\" is empty\n",
				      img, vol_name);

		warnmsg("volume size was not specified in volume \"%s\", assume "
			    "minimum to fit image \"%s\"\n", vol_name, img);
	}

	/* Fetch volume name */
	vi->name = vi_node->name;
	vi->name_len = vi_node->name_len;
	if (vi->name_len > UBI_VOL_NAME_MAX)
		return errmsg("too long volume name in section \"%s\", max. is %d characters\n",
			      vi->name, UBI_VOL_NAME_MAX);

	if (verbose)
		infomsg("volume name: %s\n", vi->name);

	/* Fetch volume alignment */
	vi->alignment = vi_node->alignment;
	if (vi->alignment == -1)
		vi->alignment = 1;
	else if (vi->id < 0)
		return errmsg("negative volume alignement %d in volume \"%s\"\n",
			      vi->alignment, vol_name);

	if (verbose)
		infomsg("volume alignment: %d\n", vi->alignment);

	/* Fetch volume flags */
	vi->flags = vi_node->flags;

	/* Initialize the rest of the volume information */
	vi->data_pad = ui->leb_size % vi->alignment;
	if (vi->data_pad != vi_node->data_pad) {
		warnmsg("the calculatd data_pad %d is not equal to "
			"the original data_pad %d!\n", vi->data_pad, vi_node->data_pad);
	}
	vi->usable_leb_size = ui->leb_size - vi->data_pad;
	if (vi->type == UBI_VID_DYNAMIC)
		vi->used_ebs = (int)((vi->bytes + vi->usable_leb_size - 1) / vi->usable_leb_size);
	else
		vi->used_ebs = (int)((st->st_size + vi->usable_leb_size - 1) / vi->usable_leb_size);
	vi->compat = 0;

	return 0;
}

int ubinize(char *out_image_file, struct volume_info_node *vi_list_head, struct ubigen_info *u, struct ubifs_info *c)
{
	struct volume_info_node *vi_node;
	struct ubigen_info ui;
	struct ubi_vtbl_record *vtbl;
	struct ubigen_vol_info *vi;
	int vols = 0;
	HANDLE h_out_img;
	/* current file pointer */
	DWORD current_offs;
	int autoresize_was_already = 0;
	char *vol_name = NULL;
	int err = -1, i;

	h_out_img = CreateFile(out_image_file, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (h_out_img == INVALID_HANDLE_VALUE) {
		errmsg("Can't create file %s: %s\n", out_image_file, last_error(msg_buf, BUF_SIZE));
		return err;
	}

	/* we don't use the subpage here, we use vid_hdr_offs only */
	ubigen_info_init(&ui, u->peb_size, u->min_io_size,
			 u->min_io_size, u->vid_hdr_offs,
			 UBI_VERSION, u->image_seq);
	
	if (verbose) {
		printf("LEB size:                  %d\n", ui.leb_size);
		printf("PEB size:                  %d\n", ui.peb_size);
		printf("min. I/O size:             %d\n", ui.min_io_size);
		//printf("sub-page size:             %d\n", subpage_size);
		printf("VID offset:                %d\n", ui.vid_hdr_offs);
		printf("data offset:               %d\n", ui.data_offs);
		printf("UBI image sequence number: %u\n", ui.image_seq);
	}

	vtbl = ubigen_create_empty_vtbl(&ui);
	if (!vtbl)
		goto out;
	
	vi_node = vi_list_head->next;
	vols = 0;
	while (vi_node != NULL) {
		vols++;
		vi_node = vi_node->next;
	}
	if (verbose)
		infomsg("count of volume: %d\n", vols);

	if (vols == 0) {
		errmsg("no volume found!\n");
		goto out_vtbl;
	}

	if (vols > ui.max_volumes) {
		errmsg("too many volumes (%d)\n", vols);
		infomsg("each vi_node corresponds to an UBI volume, maximum "
		       "count of volumes is %d\n", ui.max_volumes);
		goto out_vtbl;
	}

	vi = (struct ubigen_vol_info *)calloc(sizeof(struct ubigen_vol_info), vols);
	if (!vi) {
		errmsg("cannot allocate memory\n");
		goto out_vtbl;
	}

	/*
	 * Skip 2 PEBs at the beginning of the file for the volume table which
	 * will be written later.
	 */
	current_offs = SetFilePointer(h_out_img, ui.peb_size * 2, NULL, FILE_BEGIN);
	if (current_offs == INVALID_SET_FILE_POINTER) {
		errmsg("SetFilePointer failed: %s\n", last_error(msg_buf, BUF_SIZE));
		goto out_free;
	}
	
	if (verbose) printf("\n");

	vi_node = vi_list_head->next;
	i = 0;
	while (vi_node != NULL) {
		struct stat st;
		int j;
		HANDLE h_in_img;

		vol_name = vi_node->name;
		if (vol_name == NULL) {
			errmsg("get a volume without name\n");
			goto out_free;
		}

		if (verbose)
			infomsg("processing volume: %s\n", vol_name);

		err = read_volume_info(&ui, vi_node, &vi[i], &st);
		if (err != 0)
			goto out_free;

		if (verbose) printf("adding volume %d\n", vi[i].id);

		/*
		 * Make sure that volume ID and name is unique and that only
		 * one volume has auto-resize flag
		 */
		for (j = 0; j < i; j++) {
			if (vi[i].id == vi[j].id) {
				errmsg("volume IDs must be unique, but ID %d "
				       "in volume \"%s\" is not\n", vi[i].id, vol_name);
				goto out_free;
			}

			if (!strcmp(vi[i].name, vi[j].name)) {
				errmsg("volume name must be unique, but name \"%s\" "
					"in volume \"%s\" is not\n", vi[i].name, vol_name);
				goto out_free;
			}
		}

		if (vi[i].flags & UBI_VTBL_AUTORESIZE_FLG) {
			if (autoresize_was_already)
				return errmsg("only one volume is allowed "
					      "to have auto-resize flag\n");
			autoresize_was_already = 1;
		}

		err = ubigen_add_volume(&ui, &vi[i], vtbl);
		if (err) {
			errmsg("cannot add volume \"%s\"\n", vol_name);
			goto out_free;
		}

		if (vi_node->image_path) {
			h_in_img = CreateFile(vi_node->image_path, GENERIC_READ, FILE_SHARE_READ,
				NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
			if (h_in_img == INVALID_HANDLE_VALUE) {
				errmsg("Open %s failed!: %s\n", vi_node->image_path, last_error(msg_buf, BUF_SIZE));
				goto out_free;
			}

			if (verbose) printf("writing volume %d\n", vi[i].id);
			if (verbose) printf("image file: %s\n", vi_node->image_path);

			err = ubigen_write_volume(&ui, &vi[i], 0/* erase count */, st.st_size, h_in_img, h_out_img);
			CloseHandle(h_in_img);
			if (err) {
				errmsg("cannot write volume \"%s\"\n", vol_name);
				goto out_free;
			}
		}

		if (verbose) printf("\n");
		vi_node = vi_node->next;
	}

	if (verbose)
		infomsg("writing layout volume\n");

	err = ubigen_write_layout_vol(&ui, 0, 1, 0/* erase count */, 0/* erase count */, vtbl,h_out_img);
	if (err) {
		errmsg("cannot write layout volume\n");
		goto out_free;
	}

	if (verbose)
		infomsg("done\n");

out_free:
	free(vi);
out_vtbl:
	free(vtbl);
out:
	CloseHandle(h_out_img);
	if (err != 0)
		DeleteFile(out_image_file);

	return err;
}
